// 引用计数内联LLVM IR实现
// 纯Rust实现，无需外部C函数

use crate::{
    ast::Type,
    error::{CompilerError, Result},
    llvm_backend::LLVMBackend,
};
use inkwell::{
    types::BasicType,
    values::{BasicValueEnum, PointerValue},
    AddressSpace,
};

impl<'ctx> LLVMBackend<'ctx> {
    /// 创建引用计数块类型
    /// 结构: { strong_count: i64, weak_count: i64, data: i8* }
    pub(super) fn create_rc_block_type(&self) -> inkwell::types::StructType<'ctx> {
        let i64_type = self.context.i64_type();
        let i8_ptr_type = self.context.i8_type().ptr_type(AddressSpace::default());

        self.context.struct_type(
            &[i64_type.into(), i64_type.into(), i8_ptr_type.into()],
            false,
        )
    }

    /// 创建引用计数对象
    /// 返回指向引用计数块的指针
    pub fn generate_rc_new(
        &mut self,
        inner_type: &Type,
        _data_size: u64,
    ) -> Result<BasicValueEnum<'ctx>> {
        // 1. 创建引用计数块类型
        let rc_block_type = self.create_rc_block_type();

        // 2. 计算RC块大小并使用malloc分配
        let rc_block_size = rc_block_type.size_of().unwrap();
        let malloc_func = self
            .function_map
            .get("__malloc")
            .ok_or_else(|| CompilerError::internal("__malloc function not found"))?;
        let rc_block_ptr_i8 = self
            .builder
            .build_call(*malloc_func, &[rc_block_size.into()], "rc_block_malloc")?
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();

        // 转换为rc_block类型指针
        let rc_block_ptr = self
            .builder
            .build_bitcast(
                rc_block_ptr_i8,
                rc_block_type.ptr_type(AddressSpace::default()),
                "rc_block_ptr",
            )?
            .into_pointer_value();

        // 3. 计算数据大小并使用malloc分配
        let data_type = self.nrc_type_to_llvm_type(&Some(inner_type.clone()))?;
        let data_size_val = data_type.size_of().unwrap();
        let data_ptr_i8 = self
            .builder
            .build_call(*malloc_func, &[data_size_val.into()], "rc_data_malloc")?
            .try_as_basic_value()
            .left()
            .unwrap()
            .into_pointer_value();

        // 4. 初始化强引用计数为1
        let strong_count_ptr =
            self.builder
                .build_struct_gep(rc_block_type, rc_block_ptr, 0, "strong_count_ptr")?;
        self.builder.build_store(
            strong_count_ptr,
            self.context.i64_type().const_int(1, false),
        )?;

        // 5. 初始化弱引用计数为0
        let weak_count_ptr =
            self.builder
                .build_struct_gep(rc_block_type, rc_block_ptr, 1, "weak_count_ptr")?;
        self.builder
            .build_store(weak_count_ptr, self.context.i64_type().const_int(0, false))?;

        // 6. 设置数据指针
        let data_ptr_field =
            self.builder
                .build_struct_gep(rc_block_type, rc_block_ptr, 2, "data_ptr_field")?;
        self.builder.build_store(data_ptr_field, data_ptr_i8)?;

        Ok(rc_block_ptr.into())
    }

    /// 增加强引用计数
    pub fn generate_rc_increment(&mut self, rc_ptr: BasicValueEnum<'ctx>) -> Result<()> {
        if let BasicValueEnum::PointerValue(ptr_val) = rc_ptr {
            // 检查指针是否为nil
            let null_ptr = self
                .context
                .i8_type()
                .ptr_type(AddressSpace::default())
                .const_null();
            let is_null = self.builder.build_int_compare(
                inkwell::IntPredicate::EQ,
                ptr_val,
                null_ptr,
                "is_null",
            )?;

            // 创建条件分支
            let current_block = self.builder.get_insert_block().unwrap();
            let increment_block = self
                .context
                .append_basic_block(current_block.get_parent().unwrap(), "increment_block");
            let continue_block = self
                .context
                .append_basic_block(current_block.get_parent().unwrap(), "continue_block");

            self.builder
                .build_conditional_branch(is_null, continue_block, increment_block)?;

            // 非null分支：执行引用计数增加
            self.builder.position_at_end(increment_block);
            let rc_block_type = self.create_rc_block_type();
            let strong_count_ptr =
                self.builder
                    .build_struct_gep(rc_block_type, ptr_val, 0, "strong_count_ptr")?;
            let current_count = self.builder.build_load(
                self.context.i64_type(),
                strong_count_ptr,
                "current_count",
            )?;
            let new_count = self.builder.build_int_add(
                current_count.into_int_value(),
                self.context.i64_type().const_int(1, false),
                "new_count",
            )?;
            self.builder.build_store(strong_count_ptr, new_count)?;
            self.builder.build_unconditional_branch(continue_block)?;

            // 继续执行的代码块
            self.builder.position_at_end(continue_block);

            Ok(())
        } else {
            Err(CompilerError::internal(
                "Expected pointer value for rc_increment",
            ))
        }
    }

    /// 减少强引用计数
    pub fn generate_rc_decrement(&mut self, rc_ptr: BasicValueEnum<'ctx>) -> Result<()> {
        if let BasicValueEnum::PointerValue(ptr_val) = rc_ptr {
            // 检查指针是否为nil
            let null_ptr = self
                .context
                .i8_type()
                .ptr_type(AddressSpace::default())
                .const_null();
            let is_null = self.builder.build_int_compare(
                inkwell::IntPredicate::EQ,
                ptr_val,
                null_ptr,
                "is_null",
            )?;

            // 创建条件分支
            let current_block = self.builder.get_insert_block().unwrap();
            let decrement_block = self
                .context
                .append_basic_block(current_block.get_parent().unwrap(), "decrement_block");
            let continue_block = self
                .context
                .append_basic_block(current_block.get_parent().unwrap(), "continue_block");

            self.builder
                .build_conditional_branch(is_null, continue_block, decrement_block)?;

            // 非null分支：执行引用计数减少
            self.builder.position_at_end(decrement_block);
            let rc_block_type = self.create_rc_block_type();
            let strong_count_ptr =
                self.builder
                    .build_struct_gep(rc_block_type, ptr_val, 0, "strong_count_ptr")?;
            let current_count = self.builder.build_load(
                self.context.i64_type(),
                strong_count_ptr,
                "current_count",
            )?;
            let new_count = self.builder.build_int_sub(
                current_count.into_int_value(),
                self.context.i64_type().const_int(1, false),
                "new_count",
            )?;
            self.builder.build_store(strong_count_ptr, new_count)?;

            // 检查引用计数是否为0，如果是则释放内存
            let zero = self.context.i64_type().const_int(0, false);
            let is_zero = self.builder.build_int_compare(
                inkwell::IntPredicate::EQ,
                new_count,
                zero,
                "is_zero",
            )?;

            // 创建条件分支
            let free_block = self
                .context
                .append_basic_block(current_block.get_parent().unwrap(), "free_block");
            let decrement_continue_block = self.context.append_basic_block(
                current_block.get_parent().unwrap(),
                "decrement_continue_block",
            );

            self.builder
                .build_conditional_branch(is_zero, free_block, decrement_continue_block)?;

            // 释放内存的代码块
            self.builder.position_at_end(free_block);
            self.generate_rc_free(ptr_val)?;
            self.builder.build_unconditional_branch(continue_block)?;

            // 继续执行的代码块
            self.builder.position_at_end(decrement_continue_block);
            self.builder.build_unconditional_branch(continue_block)?;

            // 继续执行的代码块
            self.builder.position_at_end(continue_block);

            Ok(())
        } else {
            Err(CompilerError::internal(
                "Expected pointer value for rc_decrement",
            ))
        }
    }

    /// 释放引用计数对象
    fn generate_rc_free(&mut self, rc_ptr: PointerValue<'ctx>) -> Result<()> {
        let rc_block_type = self.create_rc_block_type();

        // 获取数据指针
        let data_ptr_field =
            self.builder
                .build_struct_gep(rc_block_type, rc_ptr, 2, "data_ptr_field")?;
        let data_ptr = self.builder.build_load(
            self.context.i8_type().ptr_type(AddressSpace::default()),
            data_ptr_field,
            "data_ptr",
        )?;

        // 释放数据内存
        let free_func = self
            .function_map
            .get("free")
            .ok_or_else(|| CompilerError::internal("free function not found"))?;
        self.builder
            .build_call(*free_func, &[data_ptr.into()], "free_data")?;

        // 释放引用计数块本身
        self.builder
            .build_call(*free_func, &[rc_ptr.into()], "free_rc_block")?;

        Ok(())
    }

    /// 增加弱引用计数
    pub fn generate_weak_increment(&mut self, rc_ptr: BasicValueEnum<'ctx>) -> Result<()> {
        if let BasicValueEnum::PointerValue(ptr_val) = rc_ptr {
            // 检查指针是否为nil
            let null_ptr = self
                .context
                .i8_type()
                .ptr_type(AddressSpace::default())
                .const_null();
            let is_null = self.builder.build_int_compare(
                inkwell::IntPredicate::EQ,
                ptr_val,
                null_ptr,
                "is_null",
            )?;

            // 创建条件分支
            let current_block = self.builder.get_insert_block().unwrap();
            let increment_block = self
                .context
                .append_basic_block(current_block.get_parent().unwrap(), "weak_increment_block");
            let continue_block = self
                .context
                .append_basic_block(current_block.get_parent().unwrap(), "weak_continue_block");

            self.builder
                .build_conditional_branch(is_null, continue_block, increment_block)?;

            // 非null分支：执行弱引用计数增加
            self.builder.position_at_end(increment_block);
            let rc_block_type = self.create_rc_block_type();
            let weak_count_ptr =
                self.builder
                    .build_struct_gep(rc_block_type, ptr_val, 1, "weak_count_ptr")?;
            let current_count = self.builder.build_load(
                self.context.i64_type(),
                weak_count_ptr,
                "current_count",
            )?;
            let new_count = self.builder.build_int_add(
                current_count.into_int_value(),
                self.context.i64_type().const_int(1, false),
                "new_count",
            )?;
            self.builder.build_store(weak_count_ptr, new_count)?;
            self.builder.build_unconditional_branch(continue_block)?;

            // 继续执行的代码块
            self.builder.position_at_end(continue_block);

            Ok(())
        } else {
            Err(CompilerError::internal(
                "Expected pointer value for weak_increment",
            ))
        }
    }

    /// 减少弱引用计数
    pub fn generate_weak_decrement(&mut self, rc_ptr: BasicValueEnum<'ctx>) -> Result<()> {
        if let BasicValueEnum::PointerValue(ptr_val) = rc_ptr {
            // 检查指针是否为nil
            let null_ptr = self
                .context
                .i8_type()
                .ptr_type(AddressSpace::default())
                .const_null();
            let is_null = self.builder.build_int_compare(
                inkwell::IntPredicate::EQ,
                ptr_val,
                null_ptr,
                "is_null",
            )?;

            // 创建条件分支
            let current_block = self.builder.get_insert_block().unwrap();
            let decrement_block = self
                .context
                .append_basic_block(current_block.get_parent().unwrap(), "weak_decrement_block");
            let continue_block = self
                .context
                .append_basic_block(current_block.get_parent().unwrap(), "weak_continue_block");

            self.builder
                .build_conditional_branch(is_null, continue_block, decrement_block)?;

            // 非null分支：执行弱引用计数减少
            self.builder.position_at_end(decrement_block);
            let rc_block_type = self.create_rc_block_type();
            let weak_count_ptr =
                self.builder
                    .build_struct_gep(rc_block_type, ptr_val, 1, "weak_count_ptr")?;
            let current_count = self.builder.build_load(
                self.context.i64_type(),
                weak_count_ptr,
                "current_count",
            )?;
            let new_count = self.builder.build_int_sub(
                current_count.into_int_value(),
                self.context.i64_type().const_int(1, false),
                "new_count",
            )?;
            self.builder.build_store(weak_count_ptr, new_count)?;
            self.builder.build_unconditional_branch(continue_block)?;

            // 继续执行的代码块
            self.builder.position_at_end(continue_block);

            Ok(())
        } else {
            Err(CompilerError::internal(
                "Expected pointer value for weak_decrement",
            ))
        }
    }

    /// 检查引用计数对象是否有效（强引用计数 > 0）
    pub fn generate_rc_is_valid(
        &mut self,
        rc_ptr: BasicValueEnum<'ctx>,
    ) -> Result<BasicValueEnum<'ctx>> {
        if let BasicValueEnum::PointerValue(ptr_val) = rc_ptr {
            let rc_block_type = self.create_rc_block_type();
            let strong_count_ptr =
                self.builder
                    .build_struct_gep(rc_block_type, ptr_val, 0, "strong_count_ptr")?;
            let strong_count = self.builder.build_load(
                self.context.i64_type(),
                strong_count_ptr,
                "strong_count",
            )?;
            let zero = self.context.i64_type().const_int(0, false);
            let is_valid = self.builder.build_int_compare(
                inkwell::IntPredicate::SGT,
                strong_count.into_int_value(),
                zero,
                "is_valid",
            )?;
            Ok(is_valid.into())
        } else {
            Err(CompilerError::internal(
                "Expected pointer value for rc_is_valid",
            ))
        }
    }

    /// 获取强引用计数
    pub fn generate_rc_strong_count(
        &mut self,
        rc_ptr: BasicValueEnum<'ctx>,
    ) -> Result<BasicValueEnum<'ctx>> {
        if let BasicValueEnum::PointerValue(ptr_val) = rc_ptr {
            let rc_block_type = self.create_rc_block_type();
            let strong_count_ptr =
                self.builder
                    .build_struct_gep(rc_block_type, ptr_val, 0, "strong_count_ptr")?;
            let strong_count = self.builder.build_load(
                self.context.i64_type(),
                strong_count_ptr,
                "strong_count",
            )?;
            Ok(strong_count)
        } else {
            Err(CompilerError::internal(
                "Expected pointer value for rc_strong_count",
            ))
        }
    }

    /// 获取弱引用计数
    pub fn generate_rc_weak_count(
        &mut self,
        rc_ptr: BasicValueEnum<'ctx>,
    ) -> Result<BasicValueEnum<'ctx>> {
        if let BasicValueEnum::PointerValue(ptr_val) = rc_ptr {
            let rc_block_type = self.create_rc_block_type();
            let weak_count_ptr =
                self.builder
                    .build_struct_gep(rc_block_type, ptr_val, 1, "weak_count_ptr")?;
            let weak_count =
                self.builder
                    .build_load(self.context.i64_type(), weak_count_ptr, "weak_count")?;
            Ok(weak_count)
        } else {
            Err(CompilerError::internal(
                "Expected pointer value for rc_weak_count",
            ))
        }
    }

    /// 弱引用升级为强引用
    pub fn generate_weak_lock(
        &mut self,
        rc_ptr: BasicValueEnum<'ctx>,
    ) -> Result<BasicValueEnum<'ctx>> {
        if let BasicValueEnum::PointerValue(ptr_val) = rc_ptr {
            let rc_block_type = self.create_rc_block_type();
            let rc_ptr_type = rc_block_type.ptr_type(AddressSpace::default());
            let null_ptr = rc_ptr_type.const_null();

            let current_block = self.builder.get_insert_block().unwrap();
            let parent_fn = current_block
                .get_parent()
                .ok_or_else(|| CompilerError::internal("Weak lock outside of function"))?;

            // 基本块：首先检查指针是否为null
            let check_block = self.context.append_basic_block(parent_fn, "weak_lock_check");
            let success_block = self.context.append_basic_block(parent_fn, "weak_lock_success");
            let fail_block = self.context.append_basic_block(parent_fn, "weak_lock_fail");
            let merge_block = self.context.append_basic_block(parent_fn, "weak_lock_merge");

            // entry: 先判断是否为null，null直接走fail分支
            let is_null = self.builder.build_is_null(ptr_val, "weak_lock_is_null")?;
            self.builder
                .build_conditional_branch(is_null, fail_block, check_block)?;

            // check_block: 非null时再检查强引用计数是否 > 0
            self.builder.position_at_end(check_block);
            let strong_count_ptr = self.builder.build_struct_gep(
                rc_block_type,
                ptr_val,
                0,
                "strong_count_ptr",
            )?;
            let strong_count = self.builder.build_load(
                self.context.i64_type(),
                strong_count_ptr,
                "strong_count",
            )?;
            let zero = self.context.i64_type().const_int(0, false);
            let is_valid = self.builder.build_int_compare(
                inkwell::IntPredicate::SGT,
                strong_count.into_int_value(),
                zero,
                "is_valid",
            )?;
            self.builder
                .build_conditional_branch(is_valid, success_block, fail_block)?;

            // success_block: 增加强引用计数
            self.builder.position_at_end(success_block);
            self.generate_rc_increment(ptr_val.into())?;
            let increment_exit_block = self
                .builder
                .get_insert_block()
                .ok_or_else(|| CompilerError::internal("Missing block after rc increment"))?;
            self.builder.build_unconditional_branch(merge_block)?;

            // fail_block: 直接跳转到merge_block
            self.builder.position_at_end(fail_block);
            self.builder.build_unconditional_branch(merge_block)?;

            // merge_block: 合并结果，返回非null块指针或null
            self.builder.position_at_end(merge_block);
            let phi = self.builder.build_phi(rc_ptr_type, "weak_lock_result")?;
            phi.add_incoming(&[(&ptr_val, increment_exit_block), (&null_ptr, fail_block)]);
            Ok(phi.as_basic_value().into())
        } else {
            Err(CompilerError::internal(
                "Expected pointer value for weak_lock",
            ))
        }
    }
}
